/*
 *  Copyright 2008-2009 NVIDIA Corporation
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

/*! \file monitor.h
 *  \brief Monitor iterative solver convergence
 */

#pragma once

#include <cusp/detail/config.h>

#include <cusp/blas.h>

#include <limits>
#include <iostream>
#include <iomanip>

#include <thrust/host_vector.h>

// Classes to monitor iterative solver progress, check for convergence, etc.
// Follows the implementation of Iteration in the ITL:
//   http://www.osl.iu.edu/research/itl/doc/Iteration.html

namespace cusp
{


namespace detail {

  /*! \p base_monitor : The base class for each monitor for iterative solver
   *
   *  \tparam Real real-valued type (e.g. float or double)
   *
   *  \see default_monitor
   *  \see reduction_monitor
   */
  template <class Derived, class Real>
  class base_monitor
  {
   public:
    /*! real-valued type (float or double) */
    typedef Real real_type;

    /*! \name Constructors and destructors */
    /*! \{ */

    /*! default constructor */
    base_monitor(size_t const iteration_limit,
                 Real const relative_tolerance,
                 Real const absolute_tolerance);

    /*! \} */

    /*! \name Modifiers */
    /*! \{ */
    
    /*! reset the monitor */
    void reset();

    /*! \} */

    /*! \name Iterations' control */
    /*! \{ */
    
    /*! check if this is the first iteration */
    bool first() const;
    /*! number of iterations */
    size_t iteration_count() const;
    /*! maximum number of iterations */
    size_t iteration_limit() const;
    /*! increment the iteration count */
    void operator++();

    /*! \} */

    /*! \name Tolerance's control */
    /*! \{ */

    /*! absolute tolerance */
    Real absolute_tolerance() const;
    /*! relative tolerance */
    Real relative_tolerance() const;
    /*! tolerance */
    Real tolerance() const;

    /*! \} */

    /*! \name Residue's control */
    /*! \{ */

    /*! euclidean norm of first residual */
    Real initial_residual_norm() const;
    /*! euclidean norm of last residual */
    Real residual_norm() const;
    /*! compute the residual norm from a vector */
    template <class Vector>
    void set_residual_norm(Vector const & residual);

    /*! \} */

    /*! \name Convergence's control */
    /*! \{ */

    /*! true if the last tested residual satifies the convergence tolerance */
    bool converged() const;
    /*! determine whether iteration is finished */
    template <class Vector>
    bool finished(Vector const & residual);

    /*! \} */

   protected:
    size_t const iteration_limit_;  /*!< maximum number of iterations */
    Real const relative_tolerance_; /*!< relative tolerance */
    Real const absolute_tolerance_; /*!< absolute tolerance */

    size_t iteration_count_;     /*!< number of iterations  */
    Real initial_residual_norm_; /*!< Euclidean norm of first residual */
    Real residual_norm_;         /*!< Euclidean norm of last residual */

    // noncopyable
    base_monitor(base_monitor const &);
    base_monitor & operator=(base_monitor const &);
  }; // class base monitor

  template <class Derived, class Real>
  base_monitor<Derived, Real>::base_monitor(size_t const iteration_limit,
                                            Real const relative_tolerance,
                                            Real const absolute_tolerance)
    : iteration_limit_(iteration_limit)
    , relative_tolerance_(relative_tolerance)
    , absolute_tolerance_(absolute_tolerance)
    , iteration_count_(0)
    , initial_residual_norm_(0)
    , residual_norm_(0) {}

  template <class Derived, class Real>
  void base_monitor<Derived, Real>::reset()
  {
    iteration_count_ = 0;
    initial_residual_norm_ = 0;
    residual_norm_  = 0;
  }

  template <class Derived, class Real>
  bool base_monitor<Derived, Real>::first() const
    { return (iteration_count() == 0); }

  template <class Derived, class Real>
  size_t base_monitor<Derived, Real>::iteration_count() const
    { return iteration_count_; }

  template <class Derived, class Real>
  size_t base_monitor<Derived, Real>::iteration_limit() const
    { return iteration_limit_; }

  template <class Derived, class Real>
  void base_monitor<Derived, Real>::operator++()
    { ++iteration_count_; }

  template <class Derived, class Real>
  Real base_monitor<Derived, Real>::absolute_tolerance() const
    { return absolute_tolerance_; }

  template <class Derived, class Real>
  Real base_monitor<Derived, Real>::relative_tolerance() const
    { return relative_tolerance_; }

  template <class Derived, class Real>
  Real base_monitor<Derived, Real>::tolerance() const
  {
    Derived const * derived = static_cast<Derived const *>(this);
    return derived->tolerance();
  }

  template <class Derived, class Real>
  Real base_monitor<Derived, Real>::initial_residual_norm() const
    { return initial_residual_norm_; }

  template <class Derived, class Real>
  Real base_monitor<Derived, Real>::residual_norm() const
    { return residual_norm_; }

  template <class Derived, class Real>
  template <class Vector>
  void base_monitor<Derived, Real>::set_residual_norm(Vector const & r)
  {
    residual_norm_ = cusp::blas::nrm2(r);
    if (first())
      initial_residual_norm_ = residual_norm_;
  }

  template <class Derived, class Real>
  bool base_monitor<Derived, Real>::converged() const
  {
    return (residual_norm() <= tolerance());
  }

  template <class Derived, class Real>
  template <class Vector>
  bool base_monitor<Derived, Real>::finished(Vector const & r)
  {
    set_residual_norm(r);
    return converged()
           || (iteration_count() >= iteration_limit());
  }

  /*! print out monitor information */
  template <class Derived, class Real>
  std::ostream & operator<<(std::ostream & out,
                            base_monitor<Derived, Real> const & monitor)
  {
    out << "Solver will continue until "
        << "residual norm " << monitor.tolerance() << " or reaching "
        << monitor.iteration_limit() << " iterations ";
    return out;
  }
} // end namespace detail


/*! \addtogroup iterative_solvers Iterative Solvers
 *  \addtogroup monitors Monitors
 *  \ingroup iterative_solvers
 *  \{
 */


/*! \p default_monitor : Implements standard convergence criteria
 * and reporting for iterative solvers
 *
 * \tparam Real real-valued type (e.g. \c float or \c double).
 *
 *  The following code snippet demonstrates how to configure
 *  the \p default_monitor and use it with an iterative solver.
 *
 *  \code
 *  #include <cusp/csr_matrix.h>
 *  #include <cusp/monitor.h>
 *  #include <cusp/krylov/cg.h>
 *  #include <cusp/gallery/poisson.h>
 *  
 *  int main(void)
 *  {
 *      // create an empty sparse matrix structure (CSR format)
 *      cusp::csr_matrix<int, float, cusp::device_memory> A;
 *  
 *      // initialize matrix
 *      cusp::gallery::poisson5pt(A, 10, 10);
 *  
 *      // allocate storage for solution (x) and right hand side (b)
 *      cusp::array1d<float, cusp::device_memory> x(A.num_rows, 0);
 *      cusp::array1d<float, cusp::device_memory> b(A.num_rows, 1);
 *  
 *      // set stopping criteria:
 *      //  iteration_limit    = 100
 *      //  relative_tolerance = 1e-6
 *      cusp::default_monitor<float> monitor(b, 100, 1e-6);
 *  
 *      // solve the linear system A x = b
 *      cusp::krylov::cg(A, x, b, monitor);
 *  
 *      // report solver results
 *      if (monitor.converged())
 *      {
 *          std::cout << "Solver converged to " << monitor.relative_tolerance() << " relative tolerance";
 *          std::cout << " after " << monitor.iteration_count() << " iterations" << std::endl;
 *      }
 *      else
 *      {
 *          std::cout << "Solver reached iteration limit " << monitor.iteration_limit() << " before converging";
 *          std::cout << " to " << monitor.relative_tolerance() << " relative tolerance " << std::endl;
 *      }
 *  
 *      return 0;
 *  }
 *  \endcode
 *
 *  \see \p reduction_monitor
 */
template <class Real>
class default_monitor
  : public detail::base_monitor<default_monitor<Real>, Real>
{
  typedef detail::base_monitor<default_monitor<Real>, Real> super;

 public:
  /*! \name Constructors and destructors */
  /*! \{ */

  /*! construct a \p default_monitor for a given right-hand-side \p b
   *
   *  The \p default_monitor terminates iteration when the residual norm
   *  satisfies the condition
   *       ||b - A x|| <= absolute_tolerance + relative_tolerance * ||b||
   *  or when the iteration limit is reached.
   *
   *  \param b right-hand-side of the linear system A x = b
   *  \param iteration_limit maximum number of solver iterations to allow
   *  \param relative_tolerance determines convergence criteria
   *  \param absolute_tolerance determines convergence criteria
   *
   *  \tparam VectorType vector
   */
  template <class Vector>
  default_monitor(Vector const & b,
                  size_t const iteration_limit = 500,
                  Real const relative_tolerance = 1e-5,
                  Real const absolute_tolerance = 0)
    : super(iteration_limit, relative_tolerance, absolute_tolerance)
    , b_norm_(cusp::blas::nrm2(b)) {}

  /*! \} */

  /*! \name Tolerance's control */
  /*! \{ */

  /*! \p absolute_tolerance() + \p relative_tolerance() * |b| */ 
  Real tolerance() const
  {
    return super::absolute_tolerance()
           + super::relative_tolerance() * b_norm_;
  }

  /*! \} */

 protected:
  Real const b_norm_; /*!< right-hand-size norm */
}; // class default_monitor


/*! \p reduction_monitor : Specialization of \p default_monitor which returns 
 * \c true if the initial residue is reduced by a given tolerance
 *
 * \tparam Real real-valued type (e.g. \c float or \c double).
 *
 * \see \p default_monitor
 */
template <typename Real>
class reduction_monitor
  : public detail::base_monitor<reduction_monitor<Real>, Real>
{
  typedef detail::base_monitor<reduction_monitor<Real>, Real> super;

 public:
  /*! \name Constructors and destructors */
  /*! \{ */

  /*! construct a \p reduction_monitor
   *
   *  The \p reduction_monitor terminates iteration when the residual norm
   *  satisfies the condition
   *       ||r|| <= max(absolute_tolerance, relative_tolerance * |r_0|)
   *  or when the iteration limit is reached. \c r_0 is computed the first
   *  time the function \p finished is called.
   *
   *  \param iteration_limit maximum number of solver iterations to allow
   *  \param relative_tolerance determines convergence criteria
   *  \param absolute_tolerance determines convergence criteria
   */
  reduction_monitor(size_t const iteration_limit = 500,
                    Real const relative_tolerance = 1.e-2,
                    Real const absolute_tolerance = 0)
    : super(iteration_limit,
            relative_tolerance,
            absolute_tolerance) {}

  /*! \} */

  /*! \name Tolerance's control */
  /*! \{ */

  /*! max(\p absolute_tolerance(), \p relative_tolerance() * \p initial_residual_norm()) */ 
  Real tolerance() const
  {
    return std::max(super::absolute_tolerance(),
                    super::relative_tolerance() * super::initial_residual_norm());
  }

  /*! \} */
}; // class reduction_monitor


namespace detail {

  /*! \p verbose_monitor : Displays the solver status during iteration and
   *  reports a summary after iteration has stopped
   *
   * \tparam BaseMonitor monitor used to control iterative method
   */
  template <class BaseMonitor>
  class verbose_monitor
  {
   public:
    /*! real-valued type (float or double) */
    typedef typename BaseMonitor::real_type real_type;

   public:
    /*! \name Constructors and destructors */
    /*! \{ */

    /*! construct the monitor from the output stream */
    verbose_monitor(BaseMonitor & monitor,
                    std::ostream & out)
        : monitor_(&monitor)
        , out_(out.rdbuf())
    {
        out_ << (*monitor_) << std::endl;
        out_ << "  Iteration Number  | Residual Norm" << std::endl;
    }

    /*! copy constructor */
    verbose_monitor(verbose_monitor const & monitor)
      : monitor_(monitor.monitor_)
      , out_(monitor.out_.rdbuf()) {}

    /*! \} */

    /*! \name Modifiers */
    /*! \{ */

    /*! reset the monitor  */
    void reset() const { monitor_->reset(); }

    /*! \} */

    /*! \name Iterations' control */
    /*! \{ */

    /*! check if this is the first iteration */
    bool first() const { return monitor_->first(); }
    /*! number of iterations */
    size_t iteration_count() const { return monitor_->iteration_count(); }
    /*! maximum number of iterations */
    size_t iteration_limit() const { return monitor_->iteration_limit(); }
    /*! increment the iteration count */
    void operator++() const { monitor_->operator++(); }

    /*! \} */

    /*! \name Tolerance's control */
    /*! \{ */

    /*! absolute tolerance */
    real_type absolute_tolerance() const { return monitor_->absolute_tolerance(); }
    /*! relative tolerance */
    real_type relative_tolerance() const { return monitor_->relative_tolerance(); }
    /*! tolerance */
    real_type tolerance() const { return monitor_->tolerance(); }

    /*! \} */

    /*! \name Residue's control */
    /*! \{ */

    /*! euclidean norm of first residual */
    real_type initial_residual_norm() const { return monitor_->initial_residual_norm(); }
    /*! euclidean norm of last residual */
    real_type residual_norm() const { return monitor_->residual_norm(); }
    /*! compute the residual norm from a vector */
    template <class Vector>
    void set_residual_norm(Vector const & r) const { monitor_->set_residual_norm(r); }

    /*! \} */

    /*! \name Convergence's control */
    /*! \{ */

    /*! true if the last tested residual satifies the convergence tolerance */
    bool converged() const { return monitor_->converged(); }

    /*! determine whether iteration is finished */
    template <typename Vector>
    bool finished(Vector const & r) const
    {
      bool const result = monitor_->finished(r);

      out_ << "       "  << std::setw(10) << monitor_->iteration_count();
      out_ << "       "  << std::setw(10) << std::scientific << monitor_->residual_norm() << std::endl;

      if (result)
      {
        if (monitor_->converged())
          out_ << "Successfully converged after " << monitor_->iteration_count() << " iterations." << std::endl;
        else if (monitor_->iteration_count() >= monitor_->iteration_limit())
          out_ << "Failed to converge after " << monitor_->iteration_count() << " iterations." << std::endl;
      }
      return result;
    }

    /*! \} */

   protected:
    BaseMonitor * monitor_;    /*!< the monitor used to control the iterations */
    mutable std::ostream out_; /*!< the stream used to print out the messages */
  }; // class verbose_monitor

} // end namespace detail


/*! Create a \p verbose_monitor from a generic monitor and an output stream */
template <class BaseMonitor>
detail::verbose_monitor<BaseMonitor> const
make_verbose(BaseMonitor & monitor,
             std::ostream & out = std::cout)
  { return detail::verbose_monitor<BaseMonitor>(monitor, out); }


// specialization for verbose monitor (no more verbose of verbose of ...)
template <class BaseMonitor>
detail::verbose_monitor<BaseMonitor> const &
make_verbose(detail::verbose_monitor<BaseMonitor> const & monitor,
             std::ostream & = std::cout)
  { return monitor; }


// forward declaration
namespace detail {
  template <class BaseMonitor>
  class history_monitor;
} // end namespace detail


/*! \p residual_history : Manage the history of residuals and their analysis */
template <class Real>
class residual_history
{
 public:
  /*! constructor, it reserves a number of elements equal to the iteration_limit */
  residual_history(): residuals_() {}

  /*! immediate convergence factor */
  Real immediate_rate() const
  {
    size_t const num = residuals_.size();
    if (num)
      return std::pow(residuals_[num-1] / residuals_[0], Real(1.0)/num);
    else
      return Real(0);
  }

  /*! geometric convergence factor at each step, the default value is the last step */
  Real geometric_rate(size_t step = 0) const
  {
    if (step == 0)
      step = residuals_.size() - 1;
    return residuals_[step] / residuals_[step-1];
  }

  /*! average convergence factor */
  Real average_rate() const
  {
    size_t const num = residuals_.size();
    vector average(num - 1);
    thrust::transform(residuals_.begin() + 1, residuals_.end(),
                      residuals_.begin(),
                      average.begin(),
                      thrust::divides<Real>());
    Real const sum = thrust::reduce(average.begin(), average.end(),
                                    Real(0),
                                    thrust::plus<Real>());
    return sum / (num - 1);
  }

 protected:
  typedef cusp::array1d<Real, cusp::host_memory> vector;
  vector residuals_; /*!< store the residual at each pass */

  /*! return the number of iteration */
  size_t size() const { return residuals_.size(); }

  /*! clear the history */
  void clear() { residuals_.clear(); }

  /*! reserve a memory space */
  void reserve(size_t const iteration_limit) 
    { residuals_.reserve(iteration_limit); }

  /*! the last residuals */
  Real const & back() const { return residuals_.back(); }

  /*! push back a new value */
  void push_back(Real const residual) { residuals_.push_back(residual); }

  template <class T>
  friend class detail::history_monitor;

  template <class T>
  friend std::ostream & operator<<(std::ostream & out,
                                   residual_history<T> const & history)
  {
    out << "Ran " << (history.size() - 1)
        << " iterations with a final residual of " << history.back()
        << std::endl;
    out << "geometric convergence factor : " << history.geometric_rate() << std::endl;
    out << "immediate convergence factor : " << history.immediate_rate() << std::endl;
    out << "average convergence factor : " << history.average_rate();
    return out;
  }
}; // class residual_history


namespace detail {

  /*! \p history_monitor : Save all residuals during iteration in a \p
   *  residual_history object
   *
   * \tparam BaseMonitor monitor used to control iterative method
   */
  template <class BaseMonitor>
  class history_monitor
  {
   public:
    /*! real-valued type (float or double) */
    typedef typename BaseMonitor::real_type real_type;

   public:
    /*! \name Constructors and destructors */
    /*! \{ */

    /*!/ base constructor */
    history_monitor(BaseMonitor & monitor,
                    residual_history<real_type> & history)
      : monitor_(&monitor)
      , history_(&history)
    {
      history_->clear();
      history_->reserve(monitor_->iteration_limit());
    }

    /*! copy constructor */
    history_monitor(history_monitor const & monitor)
      : monitor_(monitor.monitor_)
      , history_(monitor.history_) {}

    /*! \} */

    /*! \name Modifiers */
    /*! \{ */

    /*! reset the monitor  */
    void reset() const
    {
      monitor_->reset();
      history_->clear();
      history_->reserve(monitor_->iteration_limit());
    }

    /*! \} */

    /*! \name Iterations' control */
    /*! \{ */

    /*! check if this is the first iteration */
    bool first() const { return monitor_->first(); }
    /*! number of iterations */
    size_t iteration_count() const { return monitor_->iteration_count(); }
    /*! maximum number of iterations */
    size_t iteration_limit() const { return monitor_->iteration_limit(); }
    /*! increment the iteration count */
    void operator++() const { monitor_->operator++(); }

    /*! \} */

    /*! \name Tolerance's control */
    /*! \{ */

    /*! absolute tolerance */
    real_type absolute_tolerance() const { return monitor_->absolute_tolerance(); }
    /*! relative tolerance */
    real_type relative_tolerance() const { return monitor_->relative_tolerance(); }
    /*! tolerance */
    real_type tolerance() const { return monitor_->tolerance(); }

    /*! \} */

    /*! \name Residue's control */
    /*! \{ */

    /*! euclidean norm of first residual */
    real_type initial_residual_norm() const { return monitor_->initial_residual_norm(); }
    /*! euclidean norm of last residual */
    real_type residual_norm() const { return monitor_->residual_norm(); }
    /*! compute the residual norm from a vector */
    template <class Vector>
    void set_residual_norm(Vector const & r) const { monitor_->set_residual_norm(r); }

    /*! \} */

    /*! \name Convergence's control */
    /*! \{ */

    /*! true if the last tested residual satifies the convergence tolerance */
    bool converged() const { return monitor_->converged(); }

    /*! determine whether iteration is finished */
    template <typename Vector>
    bool finished(Vector const & r) const
    {
      bool const result = monitor_->finished(r);
      history_->push_back(monitor_->residual_norm());
      return result;
    }

    /*! \} */
   protected:
    BaseMonitor * monitor_;                 /*!< the monitor used to control the iterations */
    residual_history<real_type> * history_; /*!< the history of residuals */
  }; // class history_monitor

} // end namespace detail


/*! Create a \p history_monitor from a generic monitor and an history */
template <class BaseMonitor>
detail::history_monitor<BaseMonitor> const
save_history(BaseMonitor & monitor,
             residual_history<typename BaseMonitor::real_type> & history)
  { return detail::history_monitor<BaseMonitor>(monitor, history); }

/*! \} */

} // end namespace cusp

